import cv2
import requests
import base64
from PIL import Image
from io import BytesIO
from tqdm import tqdm

DIR = "/Users/michael/codes_ai/python/ai_demo_nlp_python/files/"
FILE_01 = "%s001.mp4" % DIR
FILE_02 = "%s002.mp4" % DIR

"""
性能优化: 使用 OpenCV 高效抽帧;tqdm 进度条可视化处理过程; 设置 API 超时防止卡死
容错机制: 异常捕获（网络错误、模型超载）; 空帧检测跳过
"""
def video_to_frames(video_path, frame_interval=1):
    """
    将视频按时间间隔抽帧
    :param video_path: 视频文件路径
    :param frame_interval: 抽帧间隔（秒）
    :return: 帧列表（PIL Image格式）
    """
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_step = int(fps * frame_interval)
    frames = []

    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % frame_step == 0:
            # 转换 OpenCV BGR 格式为 RGB
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(frame_rgb)
            frames.append(pil_image)
        frame_count += 1

    cap.release()
    return frames


def image_to_base64(pil_image):
    """
    将 PIL Image 转换为 Base64 字符串
    """
    buffered = BytesIO()
    pil_image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def analyze_video(frame, model_url="http://localhost:11434/api/generate"):
    """
    调用 Ollama API 分析单帧是否包含火灾
    """
    prompt = "请分析视频内容。"

    payload = {
        "model": "qwen2.5vl:3b",
        "prompt": prompt,
        "images": [image_to_base64(frame)],
        "stream": False,
        "options": {"temperature": 0}  # 降低随机性，确保严格判断
    }

    try:
        response = requests.post(model_url, json=payload, timeout=30000)
        if response.status_code == 200:
            result = response.json()
            return result["response"]
        else:
            print(f"API请求失败: {response.text}")
            return False
    except Exception as e:
        print(f"分析失败: {str(e)}")
        return False


def main(video_path, threshold=0.3):
    """
    主流程：分析视频火灾
    :param video_path: 视频路径
    :param threshold: 判定为火灾的帧比例阈值（默认30%）
    """
    # 抽帧处理
    print("正在提取视频帧...")
    frames = video_to_frames(video_path, frame_interval=5)  # 每n秒抽1帧
    print(f"视频帧数量: {len(frames)}")

    if not frames:
          print("视频解析失败：未提取到帧")
          return

    # 分析每帧
    print("正在分析视频内容...")
    progress_bar = tqdm(frames)

    for frame in progress_bar:
        print(analyze_video(frame))


if __name__ == "__main__":
    main(FILE_01)
    #main(FILE_02)

